#!/usr/bin/python3

import cv2
import numpy as np


def add_salt_pepper(img, SNR):
    img_ = img.copy()
    c, h, w = img_.shape
    mask = np.random.choice((0, 1, 2), size=(1, h, w), p=[SNR, (1 - SNR) / 2., (1 - SNR) / 2.])
    mask = np.repeat(mask, c, axis=0)     # 按channel 复制到 与img具有相同的shape
    img_[mask == 1] = 255    # 盐噪声
    img_[mask == 2] = 0      # 椒噪声
    return img_


img = cv2.imread('/Users/ben/Documents/AI/CSDN/pics/Lena.jpg')

img_s = add_salt_pepper(img.transpose(2, 1, 0), 0.9)
img_s = img_s.transpose(2, 1, 0)
cv2.namedWindow("salt:")
cv2.imshow('salt:', img_s)

cv2.imwrite('/Users/ben/Documents/AI/CSDN/pics/Lena_salt.jpg', img_s)

img_guassian = cv2.GaussianBlur(img, (5, 5), 0)
cv2.namedWindow("Guassian:")
cv2.imshow("Guassian:", img_guassian)

img_blur = cv2.blur(img, (5, 5))
cv2.namedWindow("Blur:")
cv2.imshow("Blur:", img_blur)

img_median = cv2.medianBlur(img_s, 5)       # 可有效去除椒盐噪声
cv2.namedWindow("medianBlur:")
cv2.imshow("medianBlur:", img_median)

ele = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
img_dilate = cv2.dilate(img, ele)
cv2.namedWindow("dilate:")
cv2.imshow("dilate:", img_dilate)

img_erode = cv2.erode(img, ele)
cv2.namedWindow("erode:")
cv2.imshow("erode:", img_erode)

img_open = cv2.morphologyEx(img, cv2.MORPH_OPEN, ele)
cv2.namedWindow("morphologyEx - open:")
cv2.imshow("morphologyEx - open:", img_open)

img_close = cv2.morphologyEx(img, cv2.MORPH_CLOSE, ele)
cv2.namedWindow("morphologyEx - close:")
cv2.imshow("morphologyEx - close:", img_close)

img_oc = cv2.morphologyEx(img_open, cv2.MORPH_CLOSE, ele)
cv2.namedWindow("morphologyEx - oc:")
cv2.imshow("morphologyEx - oc:", img_oc)

cv2.waitKey(0)

cv2.destroyAllWindows()
